import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import time

# 简单的CNN模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(32 * 7 * 7, 64)
        self.fc2 = nn.Linear(64, 10)
        self.relu = nn.ReLU()
        
    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.view(-1, 32 * 7 * 7)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

def get_device():
    """获取可用设备"""
    if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
        try:
            # 测试MPS是否真正可用
            test_tensor = torch.randn(1, 1, 2, 2).to('mps')
            test_result = test_tensor.sum()
            device = torch.device('mps')
            print("使用MPS设备")
        except:
            device = torch.device('cpu')
            print("MPS不可用，使用CPU设备")
    elif torch.cuda.is_available():
        device = torch.device('cuda')
        print("使用CUDA设备")
    else:
        device = torch.device('cpu')
        print("使用CPU设备")
    return device

def main():
    print("简单CNN示例 (支持MPS)")
    print("=" * 40)
    
    # 获取设备
    device = get_device()
    print(f"设备: {device}")
    
    # 数据预处理
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    # 加载数据
    print("加载数据...")
    trainset = torchvision.datasets.MNIST(
        root='./mnist_data', 
        train=True, 
        download=True, 
        transform=transform
    )
    
    trainloader = torch.utils.data.DataLoader(
        trainset, 
        batch_size=32,
        shuffle=True,
        num_workers=0
    )
    
    # 创建模型
    print("创建模型...")
    model = SimpleCNN()
    model.to(device)
    print(f"模型参数数量: {sum(p.numel() for p in model.parameters())}")
    
    # 损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01)
    
    # 训练一个简单的epoch
    print("开始训练...")
    model.train()
    start_time = time.time()
    
    for i, (inputs, labels) in enumerate(trainloader):
        try:
            # 数据迁移到设备
            inputs, labels = inputs.to(device), labels.to(device)
            
            # 前向传播
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            # 每100个批次打印一次
            if i % 100 == 99:
                print(f'批次 [{i+1}], 损失: {loss.item():.4f}')
                
            # 只训练前200个批次以节省时间
            if i >= 200:
                break
                
        except Exception as e:
            print(f"训练出错: {e}")
            continue
    
    training_time = time.time() - start_time
    print(f"训练完成，耗时: {training_time:.2f}秒")
    
    # 简单评估
    print("进行简单评估...")
    model.eval()
    correct = 0
    total = 0
    
    with torch.no_grad():
        for i, (inputs, labels) in enumerate(trainloader):
            try:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
                
                # 只评估前50个批次
                if i >= 50:
                    break
                    
            except Exception as e:
                print(f"评估出错: {e}")
                continue
    
    accuracy = 100 * correct / total
    print(f'准确率: {accuracy:.2f}%')
    
    print("程序执行完成!")

if __name__ == "__main__":
    main()